//
// Created by 郝羽 on 2018/1/10.
//

#ifndef RTC_XSWITCH2_AUDIO_HPP
#define RTC_XSWITCH2_AUDIO_HPP

#include <cmath>

extern "C" {
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
#include <libavutil/opt.h>
#include <libswresample/swresample.h>
}
#include <string>

#include "log.hpp"

class Audio {
public:
    std::string FilePathName() {
        return file_pathname_;
    }

    int LoadFile(const char* pathname) {
        file_pathname_ = pathname;
        av_register_all();
        avcodec_register_all();

        AVFormatContext* fmt_ctx = nullptr;
        int ret = avformat_open_input(&fmt_ctx, file_pathname_.c_str(), nullptr, nullptr);
        if (ret < 0) {
            char err[AV_ERROR_MAX_STRING_SIZE];
            av_strerror(ret, err, AV_ERROR_MAX_STRING_SIZE);
            Log::Error() << "Could not open file " << file_pathname_ << " : " << std::string(err) << std::endl;
            Reset();
            return ret;
        }
        Log::Info() << "Opened file " << file_pathname_ << std::endl;

        ret = avformat_find_stream_info(fmt_ctx, nullptr);
        if (ret < 0) {
            char err[AV_ERROR_MAX_STRING_SIZE];
            av_strerror(ret, err, AV_ERROR_MAX_STRING_SIZE);
            Log::Error() << "Could not find stream from " << file_pathname_ << " : " << std::string(err) << std::endl;
            Reset();
            return ret;
        }

        int stream_index = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_AUDIO, -1, -1, nullptr, 0);
        if (stream_index < 0) {
            Log::Error() << "Could not find audio stream from " << file_pathname_ << " : " << strError(stream_index) << std::endl;
            Reset();
            return ret;
        }

        AVStream* st = fmt_ctx->streams[stream_index];
        Log::Debug() << "Found stream " << stream_index << " codec " << std::string(avcodec_get_name(st->codecpar->codec_id)) << std::endl;
        AVCodec* dec = avcodec_find_decoder(st->codecpar->codec_id);
        if (!dec) {
            Log::Error() << "Could not find " << std::string(avcodec_get_name(st->codecpar->codec_id)) <<  " decoder" << std::endl;
            return EINVAL;
        }
        Log::Debug() << "Found " << std::string(avcodec_get_name(st->codecpar->codec_id)) << " decoder" << std::endl;

        AVCodecContext* dec_ctx = avcodec_alloc_context3(dec);
        if (!dec_ctx) {
            Log::Error() << "Could not allocate " << std::string(avcodec_get_name(st->codecpar->codec_id)) << " decoder context" << std::endl;
            return ENOMEM;
        }

        ret = avcodec_parameters_to_context(dec_ctx, st->codecpar);
        if (ret < 0) {
            Log::Error() << "Could not initialize " << std::string(avcodec_get_name(st->codecpar->codec_id)) << " decoder context : " << strError(ret) << std::endl;
            avcodec_free_context(&dec_ctx);
            return ret;
        }
        Log::Debug() << "Initialized " << std::string(avcodec_get_name(st->codecpar->codec_id)) << " decoder" << std::endl;
        AVDictionary *opts = nullptr;
        av_dict_set(&opts, "refcounted_frames", "1", 0);
        ret = avcodec_open2(dec_ctx, dec, &opts);
        if (ret < 0) {
            Log::Error() << "Could not open " << std::string(avcodec_get_name(st->codecpar->codec_id)) << " decoder context : " << strError(ret) << std::endl;
            avcodec_free_context(&dec_ctx);
            return ret;
        }

        av_dump_format(fmt_ctx, 0, file_pathname_.c_str(), 0);
        SwrContext* swr_ctx = initResampler(st);
        if (!swr_ctx) {
            Log::Error() << "Could not initialize resampler" << std::endl;
            avcodec_free_context(&dec_ctx);
            return ret;
        }

        AVCodecContext* enc_ctx = initEncoder();
        if (!enc_ctx) {
            Log::Error() << "Could not initialize encoder" << std::endl;
            avcodec_free_context(&dec_ctx);
            return ret;
        }

        AVFrame* frame_in = av_frame_alloc();
        if (!frame_in) {
            Log::Error() << "Could not allocate frame" << std::endl;
            avcodec_free_context(&dec_ctx);
            avcodec_free_context(&enc_ctx);
            return ret;
        }

        //FILE* fp = fopen("/tmp/samples.opus", "wb");
        AVPacket packet;
        while (av_read_frame(fmt_ctx, &packet) >= 0) {
            if (packet.stream_index == stream_index) {
                ret = avcodec_send_packet(dec_ctx, &packet);
                if (ret != 0 && ret != AVERROR(EAGAIN)) {
                    Log::Error() << "Could not send packet to decoder : " << strError(ret) << std::endl;
                    avcodec_free_context(&dec_ctx);
                    avcodec_free_context(&enc_ctx);
                    //av_frame_free(&frame_in);
                    av_frame_unref(frame_in);
                    return ret;
                }

                while (avcodec_receive_frame(dec_ctx, frame_in) == 0) {
                    AVFrame* frame_out = nullptr;
                    ret = doResample(swr_ctx, &frame_out, frame_in);
                    if (ret != 0) {
                        Log::Error() << "Could not resample frame" << std::endl;
                        avcodec_free_context(&dec_ctx);
                        avcodec_free_context(&enc_ctx);
                        av_frame_free(&frame_in);
                        return ret;
                    }
                    doEncode(enc_ctx, frame_out);
                    av_frame_unref(frame_out);
                }
            }
            av_packet_unref(&packet);
        }

        av_frame_free(&frame_in);

        freeResampler(&swr_ctx);

        freeEncoder(&enc_ctx);

        Log::Info() << "Loaded " << packets_.size() << " packets from " << pathname << std::endl;
        return 0;
    }

    void Reset() {
        for (auto packet : packets_) {
            av_packet_free(&packet);
        }
        packets_.clear();

        channel_layout_ = 0;
        sample_rate_ = 0;
        format_ = 0;
    }

    int Packets() {
        return packets_.size();
    }

    AVPacket* Packet(int id) {
        if (id >= packets_.size() || id < 0) {
            return nullptr;
        }
        return packets_[id];
    }

    int8_t Level(int id) {
        if (id >= levels_.size() || id < 0) {
            return -127;
        }
        return levels_[id];
    }

private:
    SwrContext* initResampler(AVStream *st) {
        SwrContext* swr_ctx = swr_alloc();
        if (!swr_ctx) {
            Log::Error() << "Could not allocate swr context" << std::endl;
            return nullptr;
        }

        if (st->codecpar->channel_layout == 0) {
            st->codecpar->channel_layout = av_get_default_channel_layout(st->codecpar->channels);
        }

        av_opt_set_channel_layout(swr_ctx, "in_channel_layout", st->codecpar->channel_layout, 0);
        av_opt_set_int(swr_ctx, "in_sample_rate", st->codecpar->sample_rate, 0);
        av_opt_set_sample_fmt(swr_ctx, "in_sample_fmt", (AVSampleFormat)st->codecpar->format, 0);
        av_opt_set_channel_layout(swr_ctx, "out_channel_layout", dst_channel_layout_, 0);
        av_opt_set_int(swr_ctx, "out_sample_rate", dst_sample_rate_, 0);
        av_opt_set_sample_fmt(swr_ctx, "out_sample_fmt", (AVSampleFormat)dst_sample_format_, 0);

        channel_layout_ = st->codecpar->channel_layout;
        sample_rate_ = st->codecpar->sample_rate;
        format_ = st->codecpar->format;

        int ret = swr_init(swr_ctx);
        if (ret < 0) {
            Log::Error() << "Could not initialize swr context : " << strError(ret) << std::endl;
            swr_free(&swr_ctx);
            return nullptr;
        }

        return swr_ctx;
    };

    void freeResampler(SwrContext** swr_ctx) {
        swr_free(swr_ctx);
    };

    int doResample(SwrContext *swr_ctx, AVFrame **dst, AVFrame *src) {
        if (src->channel_layout == 0) {
            src->channel_layout = av_get_default_channel_layout(src->channels);
        }
        if (src->channel_layout != channel_layout_ || src->sample_rate != sample_rate_ || src->format != format_) {
            Log::Error() << "Frame parameters changed" << std::endl;
            return -1;
        }

        int dst_nb_samples = av_rescale_rnd(swr_get_delay(swr_ctx, sample_rate_) + src->nb_samples, dst_sample_rate_, src->sample_rate, AV_ROUND_UP);
        //int dst_nb_samples = av_rescale(src->nb_samples, dst_sample_rate_, src->sample_rate);
//        int dst_nb_samples = dst_sample_rate_ / 1000 * dst_ptime_;
        int dst_nb_channels = av_get_channel_layout_nb_channels(dst_channel_layout_);

        AVFrame* dst2 = av_frame_alloc();
        // av_frame_unref Unreference all the buffers referenced by frame and reset the frame fields.
//        av_frame_unref(dst2);
        dst2->channel_layout = dst_channel_layout_;
        dst2->channels = dst_nb_channels;
        dst2->format = dst_sample_format_;
        dst2->nb_samples = dst_nb_samples;
        dst2->sample_rate = dst_sample_rate_;
//
        int ret = av_frame_get_buffer(dst2, 0);
        if (ret < 0) {
            Log::Error() << "Could not allocate dst frame buffer : " << strError(ret) << std::endl;
            av_frame_free(&dst2);
            return ENOMEM;
        }

        uint8_t** data = nullptr;
        int data_linesize = 0;

        int converted_samples = swr_convert(swr_ctx, dst2->data, dst_nb_samples, (const uint8_t**)src->extended_data, src->nb_samples);
        if (ret < 0) {
            Log::Error() << "Could not convert sample : " << strError(ret) << std::endl;
            av_frame_free(&dst2);
            return ret;
        }

        dst2->nb_samples = converted_samples;
        av_samples_get_buffer_size(&dst2->linesize[0], dst_nb_channels, converted_samples, dst_sample_format_, 1);

        *dst = dst2;
        Log::Debug() << "Resampled "
                     << src->nb_samples
                     << " "
                     << std::string(av_get_sample_fmt_name((AVSampleFormat)src->format))
                     << "@" << src->sample_rate
                     << " samples to "
                     << converted_samples << "/" << dst_nb_samples
                     << " "
                     << std::string(av_get_sample_fmt_name((AVSampleFormat)dst2->format))
                     << "@" << dst2->sample_rate
                     << " samples" << std::endl;
        return 0;
    };

    AVCodecContext* initEncoder() {
        AVCodec* codec = avcodec_find_encoder_by_name("libopus");
        if (!codec) {
            Log::Error() << "Could not find opus encoder" << std::endl;
            return nullptr;
        }
        AVDictionary * opts = nullptr;
        av_dict_set(&opts, "frame_duration", "20", AV_OPT_SEARCH_CHILDREN);
        av_dict_set(&opts, "application", "audio", AV_OPT_SEARCH_CHILDREN);
        av_dict_set(&opts, "compression_level", "5", AV_OPT_SEARCH_CHILDREN);
        av_dict_set(&opts, "vbr", "on", AV_OPT_SEARCH_CHILDREN);
        av_dict_set(&opts, "packet_loss", "50", AV_OPT_SEARCH_CHILDREN);
        //av_dict_set(&opts, "frame_size", "480", 0);

        AVCodecContext* enc_ctx = avcodec_alloc_context3(codec);
        if (!enc_ctx) {
            Log::Error() << "Could not allocate opus encoder context" << std::endl;
            return nullptr;
        }
        // check sample fmt
        if (codec->sample_fmts) {
            const enum AVSampleFormat* fmt = codec->sample_fmts;
            while (*fmt != AV_SAMPLE_FMT_NONE) {
                if (*fmt == dst_sample_format_) {
                    break;
                }
                ++fmt;
            }
            if (*fmt == AV_SAMPLE_FMT_NONE) {
                Log::Error() << "Opus encoder does not support sample format " << av_get_sample_fmt_name(dst_sample_format_) << std::endl;
                avcodec_free_context(&enc_ctx);
                return nullptr;
            }
        }

        // check sample rate
        if (codec->supported_samplerates) {
            const int* rate = codec->supported_samplerates;
            while (*rate != 0) {
                if (*rate == dst_sample_rate_) {
                    break;
                }
                ++rate;
            }
            if (*rate == 0) {
                Log::Error() << "Opus encoder does not support sample rate " << dst_sample_rate_ << std::endl;
                avcodec_free_context(&enc_ctx);
                return nullptr;
            }

        }

        // check channel layout
        if (codec->channel_layouts) {
            const uint64_t* channel_layout = codec->channel_layouts;
            while (*channel_layout != 0) {
                if (*channel_layout == dst_channel_layout_) {
                    break;
                }
                ++channel_layout;
            }
            if (*channel_layout == 0) {
                char layout[64] = {0,};
                av_get_channel_layout_string(layout, sizeof(layout), av_get_channel_layout_nb_channels(dst_channel_layout_), dst_channel_layout_);
                Log::Error() << "Opus encoder does not support channel layout " << std::string(layout) << std::endl;
                avcodec_free_context(&enc_ctx);
                return nullptr;
            }
        }

        enc_ctx->bit_rate = 32000;
        enc_ctx->sample_fmt = dst_sample_format_;
        enc_ctx->sample_rate = dst_sample_rate_;
        enc_ctx->time_base.den = 48000;
        enc_ctx->time_base.num = 1;

        enc_ctx->channel_layout = dst_channel_layout_;
        enc_ctx->channels = av_get_channel_layout_nb_channels(dst_channel_layout_);

        int ret = avcodec_open2(enc_ctx, codec, &opts);
        if (ret < 0) {
            Log::Error() << "Could not initialize codec opus : " << strError(ret) << std::endl;
            avcodec_free_context(&enc_ctx);
            return nullptr;
        }
        return enc_ctx;
    }


    int8_t calculateAudioLevel(int16_t *samples, size_t n) {
        if (n == 0) {
            return 0;
        }

        double rms = 0;
        for(int i = 0; i < n; i++){
            double v = static_cast<double>(samples[i]);
            v /= 32767;
            rms += v * v;
        }
        rms = std::sqrt(rms / n);
        int db = 20 * std::log10(rms);
        if (db < -127) {
            db = -127;
        } else if (db > 0) {
            db = 0;
        }
        return db;
    }

    int doEncode(AVCodecContext* enc_ctx, AVFrame* frame) {
        if (!encoding_frame_) {
            encoding_frame_ = av_frame_alloc();
            av_frame_unref(encoding_frame_);
            encoding_frame_->nb_samples = enc_ctx->frame_size;
            encoding_frame_->format = dst_sample_format_;
            encoding_frame_->sample_rate = dst_sample_rate_;
            encoding_frame_->channel_layout = dst_channel_layout_;
            encoding_frame_->channels = av_get_channel_layout_nb_channels(dst_channel_layout_);
            av_frame_get_buffer(encoding_frame_, 0);

            char layout[64] = {0,};
            av_get_channel_layout_string(layout, sizeof(layout), av_get_channel_layout_nb_channels(dst_channel_layout_), dst_channel_layout_);

            Log::Debug() << "Initialize encoding frame"
                         << " nb_samples=" << encoding_frame_->nb_samples
                         << " sample_format=" << av_get_sample_fmt_name((AVSampleFormat)encoding_frame_->format)
                         << " sample_rate=" << encoding_frame_->sample_rate
                         << " channel_layout=" << std::string(layout)
                         << std::endl;
        }
        int samples_offset = 0;
        int samples_to_copy = 0;
        int samples_remains = frame->nb_samples;
        auto level = calculateAudioLevel((int16_t*)(frame->data[0]), frame->nb_samples);
        while (samples_remains + encoding_frame_samples_ >= enc_ctx->frame_size) {
            samples_to_copy = std::min(enc_ctx->frame_size - encoding_frame_samples_, frame->nb_samples);
            av_samples_copy(encoding_frame_->data, frame->data, encoding_frame_samples_, samples_offset, samples_to_copy, av_get_channel_layout_nb_channels(dst_channel_layout_), dst_sample_format_);
            Log::Debug() << "Encoding concat " << encoding_frame_samples_ << " sample"
                         << " and " << samples_to_copy << " samples,"
                         << " remains " << samples_remains - samples_to_copy << " samples" << std::endl;
            samples_offset += samples_to_copy;
            encoding_frame_samples_ += samples_to_copy;
            samples_remains -= samples_to_copy;

            Log::Debug() << "Encoding send " << encoding_frame_samples_ << " samples" << std::endl;
            int ret = avcodec_send_frame(enc_ctx, encoding_frame_);
            encoding_frame_samples_ = 0;
            if (ret != 0 && ret != AVERROR(EAGAIN)) {
                Log::Error() << "Could not send frame to encoder : " << strError(ret) << std::endl;

                return -1;
            }

            AVPacket* packet = av_packet_alloc();
            if (!packet) {
                Log::Error() << "Could not allocate packet" << std::endl;
                avcodec_free_context(&enc_ctx);
                return -1;
            }
            //av_packet_unref(packet);
            while (avcodec_receive_packet(enc_ctx, packet) == 0) {
            //    fwrite(packet->data, 1, packet->size, fp);
                Log::Debug() << "Encoding receive packet " << packet->size << " bytes" << std::endl;
                packets_.emplace_back(packet);
                levels_.emplace_back(level);
                packet = av_packet_alloc();
                if (!packet) {
                    Log::Error() << "Could not allocate packet" << std::endl;
                    avcodec_free_context(&enc_ctx);
                    return -1;
                }
                //av_packet_unref(packet);
            }
        }

        if (samples_remains) {
            av_samples_copy(encoding_frame_->data, frame->data, 0, samples_offset, samples_remains, av_get_channel_layout_nb_channels(dst_channel_layout_), dst_sample_format_);
            encoding_frame_samples_ = samples_remains;
            Log::Debug() << "Encoding saved remaining " << samples_remains << " samples" << std::endl;
        }
        return 0;
    }

    void freeEncoder(AVCodecContext** enc_ctx) {
        avcodec_free_context(enc_ctx);
    }


    std::string strError(int ret) {
        char err[AV_ERROR_MAX_STRING_SIZE];
        av_strerror(ret, err, AV_ERROR_MAX_STRING_SIZE);
        return std::string(err);
    }
private:
    std::string file_pathname_;
    std::vector<AVPacket*> packets_;
    std::vector<int8_t> levels_;
    uint64_t channel_layout_;
    int sample_rate_;
    int format_;

    const int dst_channel_layout_ = AV_CH_LAYOUT_STEREO;
    const int dst_sample_rate_ = 48000;
    const AVSampleFormat dst_sample_format_ = AV_SAMPLE_FMT_S16;
    const int dst_ptime_ = 20;

    AVFrame* encoding_frame_ = nullptr;
    int encoding_frame_samples_ = 0;
};
#endif //RTC_XSWITCH2_AUDIO_HPP
